import os
import sys
import platform

message_register_cc=\
'''////////////////////////////////////////////////
// MACHINE GENERATED, PLEASE DONOT MODIFY!
////////////////////////////////////////////////

#include "../../plugin/protocol/protocol_register.hh"
#include "../common/generated_message_id_defines.hh"
#include "../common/generated_proto_header.hh"
#include "generated_message_handler_register.hh"

Message * ProtocolRegister::getMessage(std::uint32_t msgID) {
    auto it = msgMap_.find(msgID);
    if (it == msgMap_.end()) {
        return nullptr;
    }
    return it->second;
}

std::uint32_t ProtocolRegister::getMessageID(const Message & msg) {
    auto address = reinterpret_cast<ptrdiff_t>(msg.GetDescriptor());
    auto it = typeMap_.find(address);
    if (it == typeMap_.end()) {
        return 0;
    }
    return it->second;
}

const std::string & ProtocolRegister::getError() const {
    return error_;
}

MsgHandler * ProtocolRegister::getHandler(std::uint32_t msgID) {
    auto it = handlerMap_.find(msgID);
    if (it == handlerMap_.end()) {
        return nullptr;
    }
    return it->second;
}

'''

msg_dispatcher_hh = '''\
////////////////////////////////////////////////
// MACHINE GENERATED, PLEASE DONOT MODIFY!    //
////////////////////////////////////////////////

#pragma once

#include <cstdint>
#include <unordered_map>
#include <string>
#include "../../corelib/singleton.hpp"

class MsgDispatcher {
public:
    bool start(const std::string& cfgPath);
    std::uint16_t getSiteType(std::uint32_t msgID);
    
private:
    std::unordered_map<std::uint32_t, std::uint16_t> msgSiteMap_;
};

#define MsgDispatcherRef (*kratos::corelib::Singleton<MsgDispatcher>::instance())
'''

msg_dispatcher_cpp = '''\
////////////////////////////////////////////////
// MACHINE GENERATED, PLEASE DONOT MODIFY!    //
////////////////////////////////////////////////

#include <unordered_map>
#include <fstream>
#include <vector>
#include "msg_dispatcher.hh"
#include "../common/generated_message_id_defines.hh"
#include "../../corelib/util/string_util.hh"

bool MsgDispatcher::start(const std::string& cfgPath) {
    std::ifstream file;
    file.open(cfgPath.c_str(), std::ios::in);
    if (file.is_open()) {
        std::string line;
        while (std::getline(file, line)) {
            std::vector<std::string> result;
            kratos::corelib::StringUtil::split(line, ",", result);
            if (result.size() == 2) {
                msgSiteMap_[std::stoul(result[0])] = static_cast<std::uint16_t>(std::stoul(result[1]));
            }
        }
    } else {
        return false;
    }
    file.close();
    return true;
}

std::uint16_t MsgDispatcher::getSiteType(std::uint32_t msgID) {
    auto it = msgSiteMap_.find(msgID);
    if (it == msgSiteMap_.end()) {
        return 0;
    }
    return it->second;
}

'''

machine_generated = '''\
////////////////////////////////////////////////
// MACHINE GENERATED, PLEASE DONOT MODIFY!    //
////////////////////////////////////////////////

'''

csharp_protocol_finder = '''\
////////////////////////////////////////////////
// MACHINE GENERATED, PLEASE DONOT MODIFY!    //
////////////////////////////////////////////////

using System;
using Google.Protobuf;
using System.IO;
using System.Net;

namespace OutOfUnity
{
    public class ProtocolFinder
    {
        private MemoryStream sendStream_ = new MemoryStream(Config.SEND_BUF_SIZE);
        private NetworkMethod method_ = null;
        public ProtocolFinder(NetworkMethod method)
        {
            method_ = method;
        }
        #region PublicMethod
        public byte[] serialize(uint msgID, object o, out int length)
        {
            var pbo = o as IMessage;
            var pbSize = pbo.CalculateSize();
            var size = 8 + pbSize;
            if (size > Config.SEND_BUF_SIZE)
            {
                length = size;
                return null;
            }
            var orderedMsgID = IPAddress.HostToNetworkOrder((int)msgID);
            var orderedLength = IPAddress.HostToNetworkOrder((int)size);
            if (pbSize > 0)
            {
                sendStream_.Position = 8;
                pbo.WriteTo(sendStream_);
            }
            sendStream_.Seek(0, SeekOrigin.Begin);
            var lengthBytes = BitConverter.GetBytes(orderedLength);
            var msgIDBytes = BitConverter.GetBytes(orderedMsgID);
            sendStream_.Write(lengthBytes, 0, lengthBytes.Length);
            sendStream_.Write(msgIDBytes, 0, msgIDBytes.Length);
            sendStream_.Seek(0, SeekOrigin.Begin);
            sendStream_.Read(sendBuffer_, 0, size);
            length = size;
            return sendBuffer_;
        }
        public object deserialize(byte[] byteArray, int length, uint msgID)
        {
            var pbo = getMessage(msgID);
            pbo.MergeFrom(byteArray, 0, length);
            return pbo;
        }
        {send_method}
        #endregion
        #region PrivateMethod
        IMessage getMessage(uint msgID)
        {
            switch (msgID)
            {
                {case}
            }
            return null;
        }        
        #endregion
    }
}

'''

networking_receiver_csharp = ''''''

class MsgGenerator:
    def __init__(self):
        self.proto_files=[] # .proto files
        self.message_files = [] # .message files
        self.msg_id = 1
        self.msg_list = []
        self.msg_type = {}
        self.msg_handle = {}
        self.dispatchers = {}
        self.msg_modules = {}
        self.msg_name_id = {}
        self.msg_entity_dispatch_files = []

    def generate_doRegister_method(self):
        module_message = {}
        for file in self.message_files_list:
            module_name = file.replace('.message', '')
            for line in open('./message/' + file):
                msg = line.strip()
                module_message[msg] = module_name                
        for file in self.message_files_list:
            module_name = file.replace('.message', '')
            if not os.path.exists(module_name):
                print "Skipping: Module '" + module_name + "' not found for: " + file
                break
            output = open('./'+module_name+'/generated_protocol_register.cc', 'w')
            output.write(message_register_cc.replace('{module_name}', file.replace('.message', '')))
            output.write('bool ProtocolRegister::doRegister() {\n')
            for key, value in module_message.items():
                if value != module_name:
                    output.write('    REGISTER_MESSAGE_PROTOTYPE(' + key + ',' + key.upper() + ')\n')
            for line in open('./message/' + file):
                msg = line.strip()
                output.write('    REGISTER_MESSAGE(' + msg + ',' + msg.upper() + ')\n')
            for line in open('./message/' + file):
                msg = line.strip()
                output.write('    REGISTER_MSG_HANDLER(' + msg + ')\n')
            output.write('    return true;\n')
            output.write('}\n')
            output.close()

    def generate_protobuf_wrapper(self, file_name):
        output = open('../generated_protobuf_files/generated_' + file_name.replace('.proto', '.wrapper.cpp'), 'w')
        output.write('////////////////////////////////////////////////\n')
        output.write('// MACHINE GENERATED, PLEASE DONOT MODIFY!    //\n')
        output.write('////////////////////////////////////////////////\n\n')
        output.write('#pragma once\n\n')
        output.write('#ifdef _MSC_VER\n')
        output.write('#pragma warning(push)\n')
        output.write('#pragma warning(disable: 4100)\n')
        output.write('#pragma warning(disable: 4820)\n')
        output.write('#pragma warning(disable: 4625)\n')
        output.write('#pragma warning(disable: 4626)\n')
        output.write('#pragma warning(disable: 4514)\n')
        output.write('#pragma warning(disable: 4125)\n')
        output.write('#pragma warning(disable: 4244)\n')
        output.write('#pragma warning(disable: 4996)\n')
        output.write('#endif\n')
        output.write('#include "../protocol_files/common/protobuf_generated/' + file_name.replace('.proto', '.pb.cc"\n'))
        output.write('#ifdef _MSC_VER\n')
        output.write('#pragma warning(pop)\n')
        output.write('#endif\n')
        output.close()

    def generate_protocol_msg_id(self):
        files = os.listdir('.')
        for file in files:
            file_path = os.path.join('.', file)
            if not os.path.isdir(file_path):
                if file_path.endswith('.message'):
                    self.message_files.append(file)
                elif file.endswith('.proto'):
                    self.generate_protobuf_wrapper(file)                

    def generate_protobuf_cpp_files(self):
        files = os.listdir('./protobuf')
        for file in files:
            file_path = os.path.join('./protobuf', file)
            if not os.path.isdir(file_path):
                if file_path.endswith('.proto'):
                    self.proto_files.append(file_path)
                    os_type = platform.system()
                    if os_type =="Windows":
                        os.system('.\\protoc.exe --proto_path=.\\protobuf --cpp_out=.\\common\\protobuf_generated ' + file_path)
                        os.system('.\\protoc.exe --proto_path=.\\protobuf --python_out=.\\robot\\protobuf_generated ' + file_path)
                    else:
                        os.system('protoc --proto_path=./protobuf --cpp_out=./common/protobuf_generated ' + file_path)
                        os.system('protoc --proto_path=./protobuf --python_out=./robot/protobuf_generated ' + file_path)

    def collect_proto_files(self):
        self.proto_files_list = []
        files = os.listdir('./protobuf')
        for file in files:
            file_path = os.path.join('./protobuf', file)
            if not os.path.isdir(file_path):
                if file_path.endswith('.proto'):
                    self.proto_files_list.append(file)

    def collect_message_files(self):
        self.message_files_list = []
        files = os.listdir('./message')
        for file in files:
            file_path = os.path.join('./message', file)
            if not os.path.isdir(file_path):      
                if file_path.endswith('.message'):
                    self.message_files_list.append(file)
                elif file.endswith('.entity.dispatch'):
                    self.msg_entity_dispatch_files.append(file)

    def build_pb_cc_file_wrapper(self):
        for file_name in self.proto_files_list:
            output = open('./common/generated_' + file_name.replace('.proto', '.pb.wrapper.cc'), 'w')
            output.write('////////////////////////////////////////////////\n')
            output.write('// MACHINE GENERATED, PLEASE DONOT MODIFY!    //\n')
            output.write('////////////////////////////////////////////////\n\n')
            output.write('#ifdef _MSC_VER\n')
            output.write('#pragma warning(push)\n')
            output.write('#pragma warning(disable: 4100)\n')
            output.write('#pragma warning(disable: 4820)\n')
            output.write('#pragma warning(disable: 4625)\n')
            output.write('#pragma warning(disable: 4626)\n')
            output.write('#pragma warning(disable: 4514)\n')
            output.write('#pragma warning(disable: 4125)\n')
            output.write('#pragma warning(disable: 4244)\n')
            output.write('#pragma warning(disable: 4996)\n')
            output.write('#endif\n\n')
            output.write('#include "protobuf_generated/' + file_name.replace('.proto', '.pb.cc"\n'))
            output.write('\n#ifdef _MSC_VER\n')
            output.write('#pragma warning(pop)\n')
            output.write('#endif\n')
            output.close()

    def collect_msg_id(self):
        self.msg_id_list = []
        msg_id = 1
        for file_name in self.message_files_list:
            module_name = file_name.replace('.message', '')            
            for line in open('./message/' + file_name):
                line_strip = line.strip()
                if not self.msg_name_id.has_key(line_strip.upper()):
                    self.msg_id_list.append(line_strip)
                    self.msg_name_id[line_strip.upper()] = msg_id
                    self.msg_modules[line_strip] = module_name
                    msg_id += 1
                else:
                    print 'Error: Duplicated Message ' + line_strip
                    
    def build_entity_msg_dispatch_inl(self):
        for file_name in self.msg_entity_dispatch_files:
            module_name = file_name.replace('.entity.dispatch', '')
            file = open(os.path.join(module_name, 'entity_msg_dispatch.inl'), 'w')
            entity_dict = {}
            file.write('#include "../../plugin/base/entity.hh"\n\n')
            for line in open('./message/' + file_name):
                line_strip = line.strip()
                result = line_strip.split(' ')
                file.write('ON_ENTITY_MSG_DISPATCH(' + result[0] + ', ' + result[1] + ');\n')
                if not entity_dict.has_key(result[1]):
                    entity_dict[result[1]] = []
                entity_dict[result[1]].append(result[0])
            file.close()
            for k, v in entity_dict.items():                
                file = open(os.path.join(module_name, 'entity_msg_method_declare_' + k + '.inl'), 'w')
                file.write('public:\n')
                for msg in v:
                    file.write('DECLARE_MEMBER_MSG_HANDLER(' + msg + ');\n') 
                file.close()

    def build_msg_id_defines_hh(self):
        output = open('./common/generated_message_id_defines.hh', 'w')
        output.write('////////////////////////////////////////////////\n')
        output.write('// MACHINE GENERATED, PLEASE DONOT MODIFY!    //\n')
        output.write('////////////////////////////////////////////////\n\n')
        output.write('#pragma once\n\n')
        output.write('enum {\n')
        msg_id = 1
        for id in self.msg_id_list:
            output.write('    ' + id.upper() + '=' + str(msg_id) + ',\n')            
            msg_id += 1
        output.write('};\n')
        output.close()
        
    def build_msg_id_csharp(self):
        output = open('./csharp/ProtocolID.cs', 'w')
        output.write('////////////////////////////////////////////////\n')
        output.write('// MACHINE GENERATED, PLEASE DONOT MODIFY!    //\n')
        output.write('////////////////////////////////////////////////\n\n')
        output.write('namespace OutOfUnity\n')
        output.write('{\n')
        output.write('    class ProtocolID\n')
        output.write('    {\n')
        msg_id = 1
        for id in self.msg_id_list:
            output.write('        public const uint ' + id.upper() + '=' + str(msg_id) + ';\n')            
            msg_id += 1
        output.write('    }\n')
        output.write('}\n')
        output.close()
        
    def build_protocol_finder_csharp(self):
        output = open('./csharp/ProtocolFinder.cs', 'w')
        case_expr = ''
        for id in self.msg_id_list:
            case_expr += '            case ProtocolID.' + id.upper() + ':\n'
            case_expr += '                return new ' + id + '();\n'
        send_expr = ''
        for id in self.msg_id_list:
            send_expr += '        public bool send(' + id + ' msg)\n'
            send_expr += '        {\n'
            send_expr += '             return method_.send((int)ProtocolID.' + id.upper() + ', msg);\n'
            send_expr += '        }\n'
        output.write(csharp_protocol_finder.replace('{case}', case_expr).replace('{send_method}', send_expr))
        output.close()
        
    def build_protocol_receiver_csharp(self):
        output = open('./csharp/NetworkingReceiver.cs', 'w')
        output.write('////////////////////////////////////////////////\n')
        output.write('// MACHINE GENERATED, PLEASE DONOT MODIFY!    //\n')
        output.write('////////////////////////////////////////////////\n\n')
        output.write('using Google.Protobuf;\n')
        output.write('using UnityEngine;\n\n')
        output.write('namespace OutOfUnity\n')
        output.write('{\n')
        output.write('    public class NetworkingReceiver\n')
        output.write('    {\n')
        output.write(networking_receiver_csharp)
        for id in self.msg_id_list:
            output.write('        public delegate void OnReceive_' + id +'(' + id + ' msg);\n')
        for id in self.msg_id_list:
            output.write('        private OnReceive_' + id + ' On' + id + '_;\n')
        for id in self.msg_id_list:
            output.write('        public OnReceive_' + id + ' On' + id + '\n')
            output.write('        {\n')
            output.write('             set\n')
            output.write('             {\n')
            output.write('                 On' + id + '_ = value;\n')    
            output.write('             }\n')
            output.write('        }\n')
        output.write('    public void onMessage(uint msgID, IMessage msg)\n')
        output.write('    {\n')
        output.write('         switch (msgID)\n')
        output.write('         {\n')
        for id in self.msg_id_list:
            output.write('            case ProtocolID.' + id.upper() + ':\n')
            output.write('                if (On' + id + '_ == null)\n')
            output.write('                {\n')
            output.write('                    Debug.Log("Message handler not found:" + msg.GetType().Name);\n')
            output.write('                    return;\n')
            output.write('                }\n')
            output.write('                On' + id + '_((' + id + ')msg);\n')
            output.write('                break;\n')
        output.write('         }\n')
        output.write('    }\n')
        output.write('    }\n')
        output.write('}\n')
        output.close()

    def build_proto_wrapper_hh(self):
        output = open('./common/generated_proto_header.hh', 'w')
        output.write('////////////////////////////////////////////////\n')
        output.write('// MACHINE GENERATED, PLEASE DONOT MODIFY!    //\n')
        output.write('////////////////////////////////////////////////\n\n')
        output.write('#pragma once\n\n')
        output.write('#ifdef _MSC_VER\n')
        output.write('#pragma warning(push)\n')
        output.write('#pragma warning(disable: 4100)\n')
        output.write('#pragma warning(disable: 4820)\n')
        output.write('#pragma warning(disable: 4625)\n')
        output.write('#pragma warning(disable: 4626)\n')
        output.write('#pragma warning(disable: 4514)\n')
        output.write('#pragma warning(disable: 4125)\n')
        output.write('#pragma warning(disable: 4244)\n')
        output.write('#pragma warning(disable: 4996)\n')
        output.write('#endif\n\n')
        for file in self.proto_files_list:
            output.write('#include "protobuf_generated/' + file.replace('.proto', '.pb.h') + '"\n')
        output.write('\n#ifdef _MSC_VER\n')
        output.write('#pragma warning(pop)\n')
        output.write('#endif\n')
        output.close()

    def build_module_handler_register(self):
        for file in self.message_files_list:
            module_name = file.replace('.message', '')
            if not os.path.isdir('./' + module_name):
                os.mkdir(module_name)
            output = open('./' + module_name + '/generated_message_handler_register.hh', 'w')
            output.write('////////////////////////////////////////////////\n')
            output.write('// MACHINE GENERATED, PLEASE DONOT MODIFY!    //\n')
            output.write('////////////////////////////////////////////////\n\n')
            output.write('#pragma once\n\n')
            output.write('#include "../../plugin/protocol/protocol_register.hh"\n')
            output.write('#include "../common/generated_proto_header.hh"\n\n')
            for line in open('./message/' + file):
                name = line.strip()
                output.write('DECLARE_MSG_HANDLER(' + name + ')\n')
            output.close()
        for file in self.message_files_list:            
            module_name = file.replace('.message', '')
            output = open('./' + module_name + '/generated_module_message_handler_register.cc', 'w')
            output.write('////////////////////////////////////////////////\n')
            output.write('// MACHINE GENERATED, PLEASE DONOT MODIFY!    //\n')
            output.write('////////////////////////////////////////////////\n\n')
            output.write('#include "../../plugin/protocol/protocol_register.hh"\n')
            output.write('#include "generated_message_handler_register.hh' + '"\n\n')
            for line in open('./message/' + file):
                name = line.strip()
                output.write('DEFINE_MSG_HANDLER(' + name + ')\n')
            output.close()
            
    def collect_dispatcher(self):
        if not os.path.exists('./dispatcher/dispatcher.site'):
            return
        for line in open('./dispatcher/dispatcher.site'):
            result = line.split()
            if len(result) != 2:
                continue
            self.dispatchers[result[0].strip()] = result[1].strip()
            
    def generate_dispatcher_cfg(self):
        file = open('./dispatcher/msg_dispatcher.cfg', 'w')
        count = 1
        max_count = len(self.msg_modules)		
        for msg_name, module_id in self.msg_modules.items():
            if not self.dispatchers.has_key(module_id):
                continue
            file.write(str(self.msg_name_id[msg_name.upper()]) + ',' + self.dispatchers[module_id] + '\n')
            count += 1
        file.close()
            
    def generate_dispatcher(self):
        file_hh = open('./dispatcher/msg_dispatcher.hh', 'w')
        file_hh.write(msg_dispatcher_hh);
        file_hh.close()
        file_cpp = open('./dispatcher/msg_dispatcher.cpp', 'w')
        file_cpp.write(msg_dispatcher_cpp)
        file_cpp.close()
        
# TODO support message ID range

if __name__ == '__main__':
    generator = MsgGenerator()
    generator.collect_proto_files()
    generator.collect_message_files()
    generator.collect_msg_id()
    generator.collect_dispatcher()
    if len(sys.argv) == 2:
        if sys.argv[1] == '--update-dispatcher':
            generator.generate_dispatcher_cfg()
            generator.generate_dispatcher()
        else:
            print 'Unknown option:' + sys.argv[1]
    else:
        generator.build_msg_id_defines_hh()
        generator.build_msg_id_csharp()
        generator.build_protocol_finder_csharp()
        generator.build_protocol_receiver_csharp()
        generator.build_proto_wrapper_hh()
        generator.build_module_handler_register()
        generator.build_pb_cc_file_wrapper()
        generator.generate_doRegister_method()
        generator.generate_protobuf_cpp_files()
        generator.generate_dispatcher_cfg()
        generator.generate_dispatcher()
        generator.build_entity_msg_dispatch_inl()
